#!/usr/bin/env python3
"""
3_calculate_risk.py
风险收益计算程序
"""

import pandas as pd
import numpy as np
import sys
import os
from scipy import stats

def calculate_returns(prices):
    """计算收益率序列"""
    return prices.pct_change().dropna()

def calculate_cumulative_return(returns):
    """计算累计收益率"""
    if len(returns) == 0:
        return np.nan
    return (1 + returns).cumprod().iloc[-1] - 1

def calculate_annualized_return(returns):
    """计算年化收益率"""
    if len(returns) == 0:
        return np.nan
    return returns.mean() * 252

def calculate_max_drawdown(prices):
    """计算最大回撤"""
    if len(prices) == 0:
        return np.nan
    cumulative = (1 + prices.pct_change().fillna(0)).cumprod()
    running_max = cumulative.expanding().max()
    drawdown = (cumulative - running_max) / running_max
    return drawdown.min()

def calculate_volatility(returns, window=20):
    """计算波动率"""
    if len(returns) == 0:
        return np.nan
    return returns.rolling(window=window).std() * np.sqrt(252)

def calculate_sharpe_ratio(returns, risk_free_rate=0.03):
    """计算夏普比率"""
    if returns.std() == 0 or len(returns) == 0:
        return np.nan
    excess_returns = returns - risk_free_rate/252
    return excess_returns.mean() / excess_returns.std() * np.sqrt(252)

def calculate_sortino_ratio(returns, risk_free_rate=0.03):
    """计算索提诺比率"""
    excess_returns = returns - risk_free_rate/252
    downside_returns = excess_returns[excess_returns < 0]
    if downside_returns.std() == 0 or len(downside_returns) == 0:
        return np.nan
    return excess_returns.mean() / downside_returns.std() * np.sqrt(252)

def calculate_calmar_ratio(returns):
    """计算卡尔马比率"""
    if len(returns) == 0:
        return np.nan
    prices = (1 + returns).cumprod()
    max_dd = calculate_max_drawdown(prices)
    annual_return = calculate_annualized_return(returns)
    if max_dd == 0 or pd.isna(max_dd):
        return np.nan
    return annual_return / abs(max_dd)

def calculate_var(returns, confidence=0.05):
    """计算VaR"""
    if len(returns) == 0:
        return np.nan
    return np.percentile(returns.dropna(), confidence * 100)

def calculate_cvar(returns, confidence=0.05):
    """计算CVaR"""
    if len(returns) == 0:
        return np.nan
    var = calculate_var(returns, confidence)
    tail_losses = returns[returns <= var]
    return tail_losses.mean() if len(tail_losses) > 0 else np.nan

def calculate_beta(returns, market_returns):
    """计算Beta系数"""
    if len(returns) == 0 or len(market_returns) == 0:
        return np.nan
    if len(market_returns) != len(returns):
        min_len = min(len(returns), len(market_returns))
        returns = returns.iloc[-min_len:]
        market_returns = market_returns.iloc[-min_len:]
    
    covariance = np.cov(returns, market_returns)[0, 1]
    market_variance = np.var(market_returns)
    return covariance / market_variance if market_variance != 0 else np.nan

def calculate_tracking_error(returns, benchmark_returns):
    """计算跟踪误差"""
    if len(returns) == 0 or len(benchmark_returns) == 0:
        return np.nan
    if len(benchmark_returns) != len(returns):
        min_len = min(len(returns), len(benchmark_returns))
        returns = returns.iloc[-min_len:]
        benchmark_returns = benchmark_returns.iloc[-min_len:]
    
    active_returns = returns - benchmark_returns
    return active_returns.std() * np.sqrt(252)

def calculate_information_ratio(returns, benchmark_returns):
    """计算信息比率"""
    if len(returns) == 0 or len(benchmark_returns) == 0:
        return np.nan
    if len(benchmark_returns) != len(returns):
        min_len = min(len(returns), len(benchmark_returns))
        returns = returns.iloc[-min_len:]
        benchmark_returns = benchmark_returns.iloc[-min_len:]
    
    active_returns = returns - benchmark_returns
    tracking_error = active_returns.std() * np.sqrt(252)
    return active_returns.mean() / tracking_error * np.sqrt(252) if tracking_error != 0 else np.nan

def calculate_downside_deviation(returns, target_return=0):
    """计算下行标准差"""
    downside_returns = returns[returns < target_return/252]
    if len(downside_returns) == 0:
        return np.nan
    return downside_returns.std() * np.sqrt(252)

def calculate_upside_capture(returns, benchmark_returns):
    """计算上行捕获率"""
    if len(returns) == 0 or len(benchmark_returns) == 0:
        return np.nan
    if len(benchmark_returns) != len(returns):
        min_len = min(len(returns), len(benchmark_returns))
        returns = returns.iloc[-min_len:]
        benchmark_returns = benchmark_returns.iloc[-min_len:]
    
    up_returns = returns[benchmark_returns > 0]
    up_benchmark = benchmark_returns[benchmark_returns > 0]
    
    if len(up_returns) == 0 or up_benchmark.mean() == 0:
        return np.nan
    
    return (up_returns.mean() / up_benchmark.mean()) * 100

def calculate_downside_capture(returns, benchmark_returns):
    """计算下行捕获率"""
    if len(returns) == 0 or len(benchmark_returns) == 0:
        return np.nan
    if len(benchmark_returns) != len(returns):
        min_len = min(len(returns), len(benchmark_returns))
        returns = returns.iloc[-min_len:]
        benchmark_returns = benchmark_returns.iloc[-min_len:]
    
    down_returns = returns[benchmark_returns < 0]
    down_benchmark = benchmark_returns[benchmark_returns < 0]
    
    if len(down_returns) == 0 or down_benchmark.mean() == 0:
        return np.nan
    
    return (down_returns.mean() / down_benchmark.mean()) * 100

def calculate_skewness(returns):
    """计算偏度"""
    if len(returns) == 0:
        return np.nan
    return stats.skew(returns.dropna())

def calculate_kurtosis(returns):
    """计算峰度"""
    if len(returns) == 0:
        return np.nan
    return stats.kurtosis(returns.dropna())

def calculate_win_rate(returns):
    """计算胜率"""
    if len(returns) == 0:
        return np.nan
    return (returns > 0).sum() / len(returns.dropna())

def calculate_avg_win_loss(returns):
    """计算平均盈亏比"""
    positive_returns = returns[returns > 0]
    negative_returns = returns[returns < 0]
    
    if len(positive_returns) == 0 or len(negative_returns) == 0:
        return np.nan
    
    return positive_returns.mean() / abs(negative_returns.mean())

def calculate_risk_metrics():
    """计算所有风险收益指标"""
    
    # 读取输入数据
    input_file = '../csv_output/level2_technical.csv'
    if not os.path.exists(input_file):
        print("❌ level2_technical.csv 不存在", file=sys.stderr)
        sys.exit(1)
    
    df = pd.read_csv(input_file)
    df['date'] = pd.to_datetime(df['date'])
    df = df.sort_values('date').reset_index(drop=True)
    
    # 计算收益率
    df['daily_return'] = df['close'].pct_change()
    returns = df['daily_return'].dropna()
    
    # 计算各种风险收益指标
    df['cumulative_return'] = (1 + df['daily_return']).cumprod() - 1
    df['annualized_return'] = calculate_annualized_return(df['daily_return'])
    
    # 滚动指标
    for window in [20, 60, 252]:
        if len(df) >= window:
            df[f'volatility_{window}d'] = calculate_volatility(df['daily_return'], window)
            df[f'sharpe_{window}d'] = calculate_sharpe_ratio(df['daily_return'].iloc[-window:])
            df[f'max_drawdown_{window}d'] = calculate_max_drawdown(df['close'].iloc[-window:])
    
    # 总体指标
    df['max_drawdown_total'] = calculate_max_drawdown(df['close'])
    df['total_return'] = calculate_cumulative_return(df['daily_return'])
    df['total_volatility'] = calculate_volatility(df['daily_return'], len(df))
    
    # 风险指标
    df['var_95'] = calculate_var(returns)
    df['cvar_95'] = calculate_cvar(returns)
    df['sharpe_ratio'] = calculate_sharpe_ratio(returns)
    df['sortino_ratio'] = calculate_sortino_ratio(returns)
    df['calmar_ratio'] = calculate_calmar_ratio(returns)
    df['skewness'] = calculate_skewness(returns)
    df['kurtosis'] = calculate_kurtosis(returns)
    df['win_rate'] = calculate_win_rate(returns)
    df['avg_win_loss_ratio'] = calculate_avg_win_loss(returns)
    
    # 当前风险状态
    if len(returns) > 0:
        current_volatility = calculate_volatility(returns, min(20, len(returns)))
        df['current_risk_level'] = pd.cut([current_volatility.iloc[-1]], 
                                          bins=[0, 0.15, 0.25, 0.35, 1], 
                                          labels=['低风险', '中低风险', '中高风险', '高风险'])[0]
    else:
        df['current_risk_level'] = np.nan
    
    # 保存结果
    df.to_csv('../csv_output/level3_risk_return.csv', index=False)
    
    # 打印摘要
    print("✅ 风险收益指标计算完成")
    print(f"📊 总收益率: {df['total_return'].iloc[-1]:.2%}")
    print(f"📉 最大回撤: {df['max_drawdown_total'].iloc[-1]:.2%}")
    print(f"📈 夏普比率: {df['sharpe_ratio'].iloc[-1]:.2f}")
    print(f"⚖️  当前风险等级: {df['current_risk_level'].iloc[-1]}")

def main():
    if not os.path.exists('../csv_output/level2_technical.csv'):
        print("❌ 请先运行2_calculate_technical.py", file=sys.stderr)
        sys.exit(1)
    
    calculate_risk_metrics()
    print("✅ level3_risk_return.csv 已生成")

if __name__ == "__main__":
    main()

